package com.zhouzhou.fastjava.tree;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.zhouzhou.fastjava.util.CollectionKit;
import org.apache.commons.lang3.StringUtils;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.Function;

public class TreeHelper {

    public static <T extends TreeNode<T>, F extends FlatNode> T flat2Tree(List<F> flatNodes, Function<F, T> function) {

        List<F> rootFlatNodes = Lists.newArrayList();
        Map<String, List<F>> map = Maps.newHashMap();
        for (F flatNode : flatNodes) {
            String parentId = flatNode.getParentId();
            if (StringUtils.isNotBlank(parentId)) {
                map.computeIfAbsent(parentId, (k) -> Lists.newArrayList()).add(flatNode);
            } else {
                rootFlatNodes.add(flatNode);
            }
        }

        if (rootFlatNodes.size() != 1) {
            throw new IllegalArgumentException("Root Node Must Be Only One");
        }
        F rootFlatNode = rootFlatNodes.get(0);
        T rootTreeNode = function.apply(rootFlatNode);

        TreeHelper.buildTreeNode(rootTreeNode, map, function);

        return rootTreeNode;
    }

    public static <T extends TreeNode<T>, F extends FlatNode> List<T> flat2TreeNodes(List<F> flatNodes, Function<F, T> function) {

        List<F> rootFlatNodes = Lists.newArrayList();
        Map<String, List<F>> map = Maps.newHashMap();
        for (F flatNode : flatNodes) {
            String parentId = flatNode.getParentId();
            if (StringUtils.isNotBlank(parentId)) {
                map.computeIfAbsent(parentId, (k) -> Lists.newArrayList()).add(flatNode);
            } else {
                rootFlatNodes.add(flatNode);
            }
        }
        if (rootFlatNodes == null) {
            return null;
        }
        Set<T> roots = new TreeSet<T>();
        for (F flatNode : rootFlatNodes) {
            T rootTreeNode = function.apply(flatNode);
            TreeHelper.buildTreeNode(rootTreeNode, map, function);
            roots.add(rootTreeNode);
        }
        return Lists.newArrayList(roots);
    }


    private static <T extends TreeNode<T>, F extends FlatNode> void buildTreeNode(T treeNode, Map<String, List<F>> map, Function<F, T> function) {
        String id = treeNode.getId();
        List<F> childrenFlatNode = map.get(id);
        if (childrenFlatNode != null && !childrenFlatNode.isEmpty()) {
            for (F childFlatNode : childrenFlatNode) {
                T childTreeNode = function.apply(childFlatNode);
                TreeHelper.buildTreeNode(childTreeNode, map, function);
                treeNode.addChild(childTreeNode);
            }
        }
        boolean isExpanded = getExpanded(treeNode.getChildren());
        if (isExpanded) {
            treeNode.setExpanded(true);
        }
    }

    public static <T extends TreeNode<T>> void reBuildExpanded(List<T> nodes) {
        nodes.forEach(r -> {
            boolean isExpanded = getExpanded(r.getChildren());
            if (isExpanded) {
                r.setExpanded(true);
            }
        });
    }


    private static <T extends TreeNode<T>> boolean getExpanded(Set<T> nodes) {
        if (CollectionKit.isEmpty(nodes)) {
            return false;
        }
        boolean isExpanded = false;
        for (T node : nodes) {
            boolean b = getExpanded(node.getChildren());
            if (b) {
                node.setExpanded(true);
            }
            if (node.isExpanded()) {
                isExpanded = true;
            }

        }
        return isExpanded;
    }

    public static <T extends TreeNode<T>> List<T> getSubTreeById(List<T> nodes, final String id) {
        List<T> tagNodes = Lists.newArrayList();
        if (nodes == null) {
            return tagNodes;
        }
        nodes.forEach(r -> {
            Set<T> children = getSubTreeByRecursive(r.getChildren(), id);
            if (CollectionKit.isNotEmpty(children)) {
                tagNodes.addAll(children);
            }
            if (r.getId().equals(id)) {
                tagNodes.add(r);
            }
        });

        return tagNodes;
    }

    private static <T extends TreeNode<T>> Set<T> getSubTreeByRecursive(Set<T> nodes, final String id) {
        Set<T> set = Sets.newHashSet();
        if (nodes == null) {
            return set;
        }
        nodes.forEach(r -> {
            Set<T> children = getSubTreeByRecursive(r.getChildren(), id);
            if (CollectionKit.isNotEmpty(children)) {
                set.addAll(children);
            }

            if (r.getId().equals(id)) {
                set.add(r);
            }
        });
        return set;
    }


}
